﻿namespace X.Async
{
    using System;
    using System.Collections.Generic;
    using System.Diagnostics;
    using System.Threading;

    public class AsyncPump
    {
        public delegate void BetweenStageCall(bool hasException);

        private class AsyncPumpResult : XAsyncResult
        {
            public Exception Exp { set; get; }
            public AsyncCallback AsyncCallbackMethod { set; get; }
        }
 
        private delegate void AsyncPumpCallback(AsyncCallback cb, IAsyncResult ar);

        /// <summary>
        /// a wrapper for the callback, this class is used to save the user callback
        /// </summary>
        private class AsyncPumpCallbackShim
        {
            private readonly AsyncPumpCallback m_asyncPumperCallback;
            private readonly AsyncCallback m_userCallback;

            public AsyncPumpCallbackShim(AsyncPumpCallback asyncPumperCallback, AsyncCallback userCallback)
            {
                m_asyncPumperCallback = asyncPumperCallback;
                m_userCallback = userCallback;
            }

            public AsyncCallback EndCallback
            {
                get
                {
                    return new AsyncCallback(CallbackShim);
                }
            }

            private void CallbackShim(IAsyncResult ar)
            {
                m_asyncPumperCallback(m_userCallback, ar);
            }
        }

        public AsyncPump()
        {
        }

        public AsyncPump(StageExceptionHandlerDelegate stageExceptionHandler)
            : this()
        {
            StageExceptionHandler = stageExceptionHandler;
        }

        public object Context { private set; get; }

        public AsyncPump(StageExceptionHandlerDelegate stageExceptionHandler, object context)
            : this(stageExceptionHandler)
        {
            Context = context;
        }

        public AsyncCallback EndProc()
        {
            return EndProc(null);
        }

        public AsyncCallback EndProc(AsyncCallback cb)
        {
            IncreaseTaskCount();
            return new AsyncPumpCallbackShim(HandleEndInvoke, cb).EndCallback;
        }

        private Queue<Exception> m_expQueue = new Queue<Exception>();

        private void HandleEndInvoke(AsyncCallback cb/*this cb == the cb in EndProc*/, IAsyncResult ar)
        {
            try
            {
                if (cb != null)
                {
                    cb(ar);
                }
            }
            catch (Exception e)
            {
                //replace this by spin lock if perf here really matters
                lock (m_expQueue)
                {
                    m_expQueue.Enqueue(e);
                }
            }
            finally
            {
                DecreaseTaskCount();
            }
        }

        private IEnumerator<BetweenStageCall> m_iterator;
        private BetweenStageCall m_betweenStageCall;
        private void Pump()
        {
            //only one thread can access this function
            Debug.Assert(m_taskCount <= 0);

            try
            {
                if (m_betweenStageCall != null)
                {
                    m_betweenStageCall(m_expQueue.Count > 0);
                    m_betweenStageCall = null;

                }

                if (m_expQueue.Count > 0)
                {
                    //handle the exceptions
                    try
                    {
                        var exps = m_expQueue.ToArray();
                        if (StageExceptionHandler != null)
                        {
                            //if exception is thrown here, it means even handler cannot handle this exception
                            StageExceptionHandler(exps);
                        }
                        else
                        {
                            throw new AsyncPumpException(exps);
                        }

                    }
                    catch
                    {
                        throw;
                    }
                    finally
                    {
                        ClearStageExceptions();
                    }
                }

                // if now, m_taskCount < 0, 
                // it means CompleteInvoke has been called by either Cancel or Timeout
                if (1 == IncreaseTaskCount())
                {
                    if (m_iterator.MoveNext())
                    {
                        m_betweenStageCall = m_iterator.Current;

                        //resume pump
                        DecreaseTaskCount();

                        //from now, pump will be re-enterred
                    }
                    else
                    {
                        CompleteInvoke(null);
                    }
                }
            }
            catch (Exception e)
            {
                CompleteInvoke(e);
            }
        }

        private void CompleteInvoke(Exception exp)
        {
            // stop pumping, if original m_taskCount < 0, it means that CompleteInvoke has 
            // been called somewhere at sometime
            if (Interlocked.Exchange(ref m_taskCount, -0x19840228L) >= 0)
            {
                //m_timer may be not initialized because user specifies Timeout.Infinite
                if (m_timer != null)
                {
                    m_timer.Dispose();
                    m_timer = null;
                }

                m_ar.Exp = exp;
                m_ar.IsCompleted = true; //it will signal the event inside ar

                if (m_ar.AsyncCallbackMethod != null)
                {
                    m_ar.AsyncCallbackMethod(m_ar);
                }
            }
        }

        public delegate void StageExceptionHandlerDelegate(Exception[] exps);
        private StageExceptionHandlerDelegate StageExceptionHandler { set; get; }

        private long m_taskCount = 0;

        private long IncreaseTaskCount()
        {
            return Interlocked.Increment(ref m_taskCount);
        }

        private void DecreaseTaskCount()
        {
            if (0 == Interlocked.Decrement(ref m_taskCount))
            {
                // go to next stage
                Pump();
            }
        }

        public Exception[] GetStageExceptions()
        {
            Debug.Assert(Interlocked.Read(ref m_taskCount) == 0, "this method should only be called between stages");
            return m_expQueue.ToArray();
        }

        public void ClearStageExceptions()
        {
            Debug.Assert(Interlocked.Read(ref m_taskCount) == 0, "this method should only be called between stages");
            m_expQueue.Clear();
        }


        private AsyncPumpResult m_ar;
        private Timer m_timer;

        private void HandleTimeout(object state)
        {
            CompleteInvoke(new TimeoutException());
        }


        private void CancelNextStageOperation()
        {
            CompleteInvoke(new OperationCanceledException());
        }


        public IAsyncResult BeginInvoke(IEnumerator<BetweenStageCall> iterator, AsyncCallback callback, object state, int timeout)
        {
            Debug.Assert(m_ar == null);
            m_ar = new AsyncPumpResult();
            m_ar.AsyncCallbackMethod = callback;
            m_ar.AsyncState = state;
            m_iterator = iterator;
            Pump();
            if (timeout != Timeout.Infinite)
            {
                m_timer = new Timer(HandleTimeout, null, timeout, 0);
            }
            return m_ar;
            
        }

        public IAsyncResult BeginInvoke(IEnumerator<BetweenStageCall> iterator, AsyncCallback callback, object state)
        {
            return BeginInvoke(iterator, callback, state, Timeout.Infinite);
        }


        public void EndInvoke(IAsyncResult ar)
        {
            try
            {
                AsyncPumpResult result = (AsyncPumpResult)ar;

                if (!result.IsCompleted)
                {
                    ar.AsyncWaitHandle.WaitOne();
                }

                if (m_ar.Exp != null)
                    throw m_ar.Exp;
            }
            finally
            {
                //clean up the resource
                m_iterator = null;
                m_expQueue = null;
                m_ar = null;
                m_betweenStageCall = null;
                Debug.Assert(m_timer == null);
                //don't touch m_taskCount here, we still need it for other callbacks
            }
        }


        public void Invoke(IEnumerator<BetweenStageCall> iterator, int timeout)
        {
            EndInvoke(
                BeginInvoke(iterator, null, null, timeout));
        }


        public void Invoke(IEnumerator<BetweenStageCall> iterator)
        {
            Invoke(iterator, Timeout.Infinite);
        }

    }

    public sealed class AsyncPump<TResult> : AsyncPump
    {
        public AsyncPump()
            : base()
        {
        }

        public AsyncPump(StageExceptionHandlerDelegate stageExceptionHandler)
            : base(stageExceptionHandler)
        {
        }

        public AsyncPump(StageExceptionHandlerDelegate stageExceptionHandler, object context)
            : base(stageExceptionHandler, context)
        {
        }

        public TResult Result { set; private get; }

        public new TResult Invoke(IEnumerator<BetweenStageCall> iterator, int timeout)
        {
            EndInvoke(
                BeginInvoke(iterator, null, null, timeout));

            return Result;
        }

        public new TResult Invoke(IEnumerator<BetweenStageCall> iterator)
        {
            return Invoke(iterator, Timeout.Infinite);
        }

        public new TResult EndInvoke(IAsyncResult ar)
        {
            base.EndInvoke(ar);
            return Result;
        }

    }
}
